
import cv2
from .imageProcess import watershed
from .flafeaqty import MeaArea



def get_data_label(dataset):
    DataList = []
    LabelList = []
    for iLine in dataset:
        iLineData = iLine.rstrip().split()
        img_path = iLineData[0]
        image = cv2.imread(img_path)
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        binary = watershed(image)
        frame_mean, frame_mea, frame_area = MeaArea(gray, binary)
        iDataList = [frame_mean, frame_mea, frame_area]
        DataList.append(iDataList)
        LabelList.append(int(iLineData[1]))

    return DataList, LabelList


def get_data(dataset):
    DataList = []
    for iLine in dataset:
        iLineData = iLine.rstrip().split()
        img_path = iLineData[0]
        image = cv2.imread(img_path)
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        binary = watershed(image)
        frame_mean, frame_mea, frame_area = MeaArea(gray, binary)
        iDataList = [frame_mean, frame_mea, frame_area, int(iLineData[1])]
        DataList.append(iDataList)

    return DataList

